#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <vector>
#include <iostream>
#include <numeric>
namespace bst = boost::numeric::ublas;

//precondition: p & q are initialized, w is precomputed
//[i,j)
int optimal_bst_bottomup(const std::vector<double> & p, const std::vector<double> & q,
						bst::matrix<double> & e, bst::matrix<double> & w, bst::matrix<int> & root) 
{
	e.resize(q.size(),q.size());
	root.resize(q.size(),q.size());
	for ( int i = 0; i < q.size(); ++i)
		e(i,i) = q[i];
	for (int i = p.size()-1; i >= 0; --i) {
		for (int j = i+1; j <= p.size(); ++j) {
			e(i,j) = std::numeric_limits<double>::max();
			for (int k = i; k < j; ++k) {
				double tmp = e(i,k) + e(k+1,j) + w(i,j);
				if (e(i,j) > tmp) {
					e(i,j) = tmp;
					root(i,j) = k;
				}
			}
		}	
	}
	return e(0, p.size());
}

void compute_w(const std::vector<double> & p, const std::vector<double> & q, bst::matrix<double> & w) {
	w.resize(p.size(),q.size());
	for ( int i = 0; i < p.size(); ++i) {
		w(i,i) = q[i];
		for( int j = i + 1; j <= p.size(); ++j) {
			w(i,j) = w(i,j-1) + p[j-1] + q[j];
		}
	}
}

//call [i,j)
void print_optimal_bst(const bst::matrix<int> & root, int i, int j){
	if(i == j){
		std::cout << "D" << i << ' ';
	}
	else {
		int k = root(i,j);
		std::cout << "K" << k << ' ';
		print_optimal_bst(root, i, k);
		print_optimal_bst(root, k+1,j);
	}
}

void test_opt_bst(double * pa, int m, double * qa, int n){
	std::vector<double> p(pa, pa+m);
	std::vector<double> q(qa, qa+n);
	bst::matrix<double> e,w;
	bst::matrix<int> root;
	
	compute_w(p,q,w);
	std::cout << "matrix w:\n" << w << std::endl;
	optimal_bst_bottomup(p,q,e,w,root);
	std::cout << "matrix e:\n" << e << std::endl;
	std::cout << "matrix root:\n" <<root << std::endl;
	std::cout << "expected search cost: " << e(0,p.size()) << std::endl;
	print_optimal_bst(root,0,p.size()); std::cout << std::endl;
}

void test_optimal_bst(){
	double pa[] = {0.04, 0.06, 0.08, 0.02, 0.10, 0.12, 0.14};
	double qa[] = {0.06, 0.06, 0.06, 0.06, 0.05, 0.05, 0.05, 0.05};
	test_opt_bst(pa,sizeof(pa)/sizeof(double), qa, sizeof(qa)/sizeof(double));
	
	double pb[] = {0.15, 0.10, 0.05, 0.10, 0.20};
	double qb[] = {0.05, 0.10, 0.05, 0.05, 0.05, 0.10};
	test_opt_bst(pb, sizeof(pb)/sizeof(double), qb, sizeof (qb)/sizeof(double));
}